'''
Created on 25-mei-2013

@author: Michael
'''

# http://stackoverflow.com/questions/11294859/how-to-define-the-markers-for-watershed-in-opencv


from visualization import drawscaled
from preprocessing import opening, closing, adaptiveThreshold
from imageReader import readCoords, writeImage, getPanoramicImage, writeReportImage
import cv2
import numpy as np

from scipy.ndimage import label

def preprocessForWaterShed(img):
    image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    edgemap = closing(opening(adaptiveThreshold(image)))
    return edgemap


    
    
def segment_on_dt(a, img):
    border = cv2.dilate(img, None, iterations=5)
    border = border - cv2.erode(border, None)

    dt = cv2.distanceTransform(img, 2, 3)
    dt = ((dt - dt.min()) / (dt.max() - dt.min()) * 255).astype(np.uint8)
    _, dt = cv2.threshold(dt, 45, 255, cv2.THRESH_BINARY)
    lbl, ncc = label(dt)
    lbl = lbl * (255/ncc)
    # Completing the markers now. 
    lbl[border == 255] = 255

    lbl = lbl.astype(np.int32)
    cv2.watershed(a, lbl)

    lbl[lbl == -1] = 0
    lbl = lbl.astype(np.uint8)
    return 255 - lbl

def watershed(img):
    img_gray = preprocessForWaterShed(img)
    image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    for rowIndex, row in enumerate(img_gray):
        for columnIndex, value in enumerate(row):
            image[rowIndex,columnIndex] = max(image[rowIndex,columnIndex] - value*0.45,0)
    _, img_bin = cv2.threshold(image, 0, 255, cv2.THRESH_OTSU)
    img_bin = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, np.ones((3, 3), dtype=int))
    result = segment_on_dt(img, img_bin)
    return result

def getPrimaryArea(image):
    watershedded = watershed(image)
    middleColumn = int(watershedded.shape[1]/2)
    countMap = dict()
    for i in range(0,256):
        countMap[i] = 0
    for v in watershedded[:,middleColumn]:
        countMap[v] += 1
    primaryColor = 0
    primaryCount = 0
    for key in countMap.keys():
        if(countMap[key] > primaryCount):
            primaryColor = key
            primaryCount = countMap[key]
    watershedded[watershedded != primaryColor] = 0
    watershedded[watershedded == primaryColor] = 100
    return watershedded

def getHighlightedTopTeeth(image,teethImage,teethImageNr):
    teethImage = getPrimaryArea(teethImage)
    newImage = np.zeros((image.shape[0],image.shape[1]),np.uint8)
    (translateRow,translateColumn) = readCoords("Teeth_coordinates/topteeth",teethImageNr)
    for rowIndex,row in enumerate(teethImage):
        for columnIndex,value in enumerate(row):
            newImage[rowIndex+translateRow,columnIndex+translateColumn] = value
    cv2.imwrite("Watershedded images/top_teeth_" + str(teethImageNr) + ".png", newImage)
    
def getHighlightedBottomTeeth(image,teethImage,teethImageNr):
    teethImage = getPrimaryArea(teethImage)
    newImage = np.zeros((image.shape[0],image.shape[1]),np.uint8)
    (translateRow,translateColumn) = readCoords("Teeth_coordinates/bottomteeth",teethImageNr)
    for rowIndex,row in enumerate(teethImage):
        for columnIndex,value in enumerate(row):
            newImage[rowIndex+translateRow,columnIndex+translateColumn] = value
    cv2.imwrite("Watershedded images/bottom_teeth_" + str(teethImageNr) + ".png", newImage)
    
